package com.cook.book.controller;

import cn.hutool.core.date.DatePattern;
import cn.hutool.core.date.DateUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.util.Date;
import java.util.concurrent.CompletableFuture;

/** 处理大模型流式响应 */
@Slf4j
@RestController
@RequestMapping("sse")
public class SseController {

    @PostMapping(value = "/emitter", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter emitter() {
        SseEmitter sseEmitter = new SseEmitter();
        // 异步线程处理，避免阻塞
        CompletableFuture.runAsync(() -> {
            for (int i = 0; i < 10; i++) {
                try {
                    Thread.sleep(2000); // 模拟延迟
                    sseEmitter.send("数据生成时间: " + DateUtil.format(new Date(), DatePattern.NORM_DATETIME_PATTERN));
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
            }
            sseEmitter.complete();
        });
        return sseEmitter;
    }

    @PostMapping(value = "/webflux", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> webflux() {
        // 异步线程处理，避免阻塞
        return Flux.create(sink -> CompletableFuture.runAsync(() -> {
            for (int i = 0; i < 10; i++) {
                try {
                    Thread.sleep(2000); // 模拟延迟
                    sink.next("数据生成时间: " + DateUtil.format(new Date(), DatePattern.NORM_DATETIME_PATTERN));
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
            }
            sink.complete();
        }));
    }

    // 依次处理并返回多种情况的结果
    @PostMapping(value = "/more", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> more() {
        // SecurityContext context = SecurityContextHolder.getContext();
        return Flux.create(sink -> {
            CompletableFuture<Void> future = CompletableFuture.completedFuture(null);

            for (int i = 0; i < 5; i++) {
                int finalI = i;
                future = future.thenCompose(v -> CompletableFuture.runAsync(() -> {
                    // 涉及用户信息，SpringSecurity同步上下文
                    // SecurityContextHolder.setContext(context);
                    for (int j = 0; j < 10; j++) {
                        try {
                            Thread.sleep(2000); // 模拟延迟
                            sink.next("数据生成时间: " + DateUtil.format(new Date(), DatePattern.NORM_DATETIME_PATTERN));
                            log.info("test {} {}", finalI, j);
                        } catch (Exception e) {
                            throw new RuntimeException(e);
                        }
                    }
                }));
            }

            future.thenRun(sink::complete).exceptionally(ex -> {
                sink.error(ex);
                return null;
            });
        });
    }

}
